import os
import argparse
import ast

import tensorflow as tf  # 导入tensorflow库
import moxing as mox
import security_utils


parser = argparse.ArgumentParser(description='Image classification')
parser.add_argument('--data_url',
                    metavar='DIR',
                    default='/cache/data_url',
                    help='path to dataset')
parser.add_argument('--train_url',
                    default="/mindspore-dataset/output/",
                    type=str,
                    help="setting dir of training output")

parser.add_argument('--data_encrypted', type=ast.literal_eval, default=False,
                    help="data is encrypted, "
                         "need to be decrypted before training")

parser.add_argument('--encrypt_model', type=ast.literal_eval, default=False,
                    help="model file need to be encrypted after training")

parser.add_argument('--user_dir_prefix', type=str, default='/cache',
                    choices=("/cache", "/home/ma-user/work/cache"),
                    help="user directory prefix")

args_opt = parser.parse_args()

USER_DIR_PREFIX = args_opt.user_dir_prefix

CUR_RANK_ID = os.environ.get('RANK_ID')
CACHE_TRAINING_URL = os.path.join(USER_DIR_PREFIX, f'training/{CUR_RANK_ID}/')

if not os.path.isdir(CACHE_TRAINING_URL):
    os.makedirs(CACHE_TRAINING_URL)


if __name__ == '__main__':
    real_path = os.path.join(USER_DIR_PREFIX, 'data_url/0')
    mox.file.copy_parallel(args_opt.data_url, real_path)
    real_data_parent_path = real_path
    dew_config_file_path = os.path.join(real_data_parent_path, 'dew.ini')
    cipher_data_dir = os.path.join(real_data_parent_path, 'data_cipher')
    if args_opt.data_encrypted:
        print(f'[{CUR_RANK_ID}] begin to get plain data...')
        plain_data_dir = security_utils.decrypt_dir(dew_config_file_path,
                                                    cipher_data_dir,
                                                    USER_DIR_PREFIX,
                                                    'decrypt_data.flag')
        print(f'[{CUR_RANK_ID}] get plain data successfully')
    else:
        plain_data_dir = cipher_data_dir
    real_path = plain_data_dir
    os.system(f"mkdir -p ~/.keras/datasets/&&cp -rf {plain_data_dir}/* ~/.keras/datasets/")

    # 加载数据
    mnist = tf.keras.datasets.mnist
    (x_train, y_train), (x_test, y_test) = mnist.load_data()

    # 搭建模型
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax))
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    # 数据预处理
    x_train, y_train = x_train[:1000] / 255, y_train[:1000]
    x_test, y_test = x_test[:1000] / 255, y_test[:1000]

    # 训练
    model.fit(x_train, y_train, epochs=5)

    # 评估
    val_loss, val_acc = model.evaluate(x_test, y_test)
    print('first evaluate loss: {} acc: {}'.format(val_loss, val_acc))

    # 加密保存模型
    security_utils.encrypt_model(dew_config_file_path, model, CACHE_TRAINING_URL)
    # 上传到OBS
    mox.file.copy_parallel(CACHE_TRAINING_URL, args_opt.train_url)
